// Copyright (c) 2019 Presto Labs Pte. Ltd.
// Author: donggu

#pragma once
#include <algorithm>
#include <ctime>  // strptime
#include <experimental/filesystem>
#include <map>
#include <memory>
#include <regex>
#include <set>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include <limits>

// external
#include <boost/algorithm/string/predicate.hpp>
#include <fmt/format.h>

// internal
#include "coin2/exchange/base/executor/topic_publisher.h"
#include "orio/io/base.h"
#include "orio/io/concat.h"
#include "orio/record/topic.h"

namespace coin2::flow {

struct FlowFile {
  std::experimental::filesystem::path path;
  std::string topic_string;
  int64_t timestamp;
};

inline bool operator<(const FlowFile& lhs, const FlowFile& rhs) {
  return std::tie(lhs.timestamp, lhs.topic_string, lhs.path) <
         std::tie(rhs.timestamp, rhs.topic_string, rhs.path);
}

using TopicRecordReader = orio::io::Reader<orio::record::TopicRecord>;

class TopicIdInjector : public TopicRecordReader {
 public:
  TopicIdInjector(uint64_t topic_id, std::unique_ptr<TopicRecordReader> reader)
      : topic_id_(topic_id), reader_(std::move(reader)) {}

  bool Read(orio::record::TopicRecord* out) {
    if (!reader_->Read(out)) {
      return false;
    }
    out->topic_id = topic_id_;
    return true;
  }

 private:
  uint64_t topic_id_;
  std::unique_ptr<TopicRecordReader> reader_;
};

class ShiftTopicReader : public TopicRecordReader {
 public:
  ShiftTopicReader(std::vector<std::unique_ptr<orio::record::MergedTopicReader>> readers,
                   std::vector<orio::record::TopicRecord*> buffers)
      : readers_(std::move(readers)), buffers_(std::move(buffers)) {
    CHECK_GT(readers_.size(), 0);
  }

  bool Read(orio::record::TopicRecord* out) override {
    if (primary_reader_ >= 0 && readers_[primary_reader_]) {
      if (readers_[primary_reader_]->Read(out)) {
        int64_t diff = out->timestamp - last_timstamp_;
        if (diff < threshold_) {
          last_timstamp_ = out->timestamp;
          return true;
        } else {
          // the reader may switch, don't drop already read data, copy data to buffer
          memcpy(reinterpret_cast<void*>(buffers_[primary_reader_]),
              reinterpret_cast<const void*>(out), out->struct_size());
        }
      } else {
        readers_[primary_reader_].reset();
      }
    }

    // primary reader gap or no primary reader, need switch
    if (Switch()) {
      // the lastest message already read into buffer after switch
      memcpy(reinterpret_cast<void*>(out), reinterpret_cast<const void*>(buffers_[primary_reader_]),
          buffers_[primary_reader_]->struct_size());
      last_timstamp_ = out->timestamp;
      return true;
    }

    return false;
  }

 private:
  bool Switch() {
    // read first message, over the last timestamp
    primary_reader_ = ReadOver();
    return primary_reader_ >= 0;
  }

  // read unit the the first of bigger than last timestamp
  int ReadOver() {
    int64_t min_timestamp = std::numeric_limits<int64_t>::max();
    int index = -1;
    if (primary_reader_ >= 0 && readers_[primary_reader_]) {
      // primary_reader is still valid
      min_timestamp = buffers_[primary_reader_]->timestamp;
      index = primary_reader_;
    }

    for (int i = 0; i < readers_.size(); i++) {
      // skip the primary reader as it is already over last timestamp
      if (i == primary_reader_ || !readers_[i]) {
        continue;
      }

      bool status = false;
      while ((status = readers_[i]->Read(buffers_[i]))) {
        if (buffers_[i]->timestamp > last_timstamp_) {
          break;
        }
      }

      if (status) {
        if (buffers_[i]->timestamp < min_timestamp) {
          min_timestamp = buffers_[i]->timestamp;
          index = i;
        }
      } else {
        readers_[i].reset();
      }
    }

    return index;
  }

  std::vector<std::unique_ptr<orio::record::MergedTopicReader>> readers_;
  std::vector<orio::record::TopicRecord*> buffers_;
  int64_t last_timstamp_ = 0;
  int primary_reader_ = -1;
  static constexpr int64_t threshold_ = 10 * 1000000000LL;  // 10 seconds
};

// We read orio::io::TopicRecord (w/ topic_id=0)
// and convert it to TopicRecordPacket (w/ topic string)
class FlowReader : public orio::io::Reader<exchange::base::executor::TopicRecordPacket> {
 public:
  FlowReader(const std::string& machine, const std::string& root_dir,
             const std::set<std::string>& topics, int64_t begin_timestamp, int64_t end_timestamp,
             size_t max_record_size, bool shift)
    : max_record_size_(max_record_size),
      buf_(new char[max_record_size]),
      topics_(std::vector<std::string>(std::begin(topics), std::end(topics))),
      shift_(shift) {
    std::vector<std::unique_ptr<TopicRecordReader>> readers;
    std::vector<std::string> topic_strings;

    // reader1 = IntervalTopicReader(
    //  ConcatReader(([TopicReader(FileReader(topic1-day1)), TopicReader(FileReader(topic1-day2),
    //  ..])), from, to
    // )
    //
    // reader2 = IntervalTopicReader(
    //  ..
    //  from, to
    // )
    // ..
    // return MergedTopicReader(reader1, reader2, ..)

    // get file list
    const std::vector<FlowFile> files =
        GetFileList(machine, root_dir, begin_timestamp, end_timestamp);

    std::vector<orio::record::TopicRecord*> buffers;
    for (int topic_id = 0; topic_id < topics_.size(); topic_id++) {
      const std::string& topic = topics_[topic_id];
      std::vector<std::experimental::filesystem::path> sub_paths;

      // filter with topic
      for (const auto& file : files) {
        // topic =
        // 'binance_ws_trade'
        // 'binance_ws_trade_1'
        // 'binance_ws_trade_a1'
        // /remote/iosg/coin/data/flow/{machine}/{date}/files
        if (file.topic_string == topic) {
          sub_paths.push_back(file.path);
        }
      }
      if (sub_paths.empty()) {
        continue;
      }
      LOG(INFO) << "Read raw files for topic: " << topic;
      for (const auto& sub_path : sub_paths) {
        LOG(INFO) << sub_path;
      }
      LOG(INFO) << "\n";

      buffers_.push_back(std::make_unique<char[]>(max_record_size));
      buffers.push_back(reinterpret_cast<orio::record::TopicRecord*>(buffers_.back().get()));

      std::unique_ptr<TopicRecordReader> concat_topic_reader(
          new orio::record::ConcatTopicReader(std::move(sub_paths), max_record_size, false));

      readers.push_back(std::unique_ptr<TopicRecordReader>(new TopicIdInjector(
          topic_id, std::make_unique<orio::record::IntervalTopicReader>(
            std::move(concat_topic_reader), begin_timestamp, end_timestamp))));
      topic_strings.push_back(topic);
    }
    if (readers.size() == 1) {
      reader_ = std::move(readers[0]);
    } else if (readers.size() > 1) {
      if (shift_) {
        std::map<char, std::vector<std::unique_ptr<TopicRecordReader>>> worker_to_readers;
        for (int i = 0; i < readers.size(); ++i) {
          char worker = topic_strings[i][topic_strings[i].size() - 1];
          if (worker_to_readers.count(worker) == 0) {
            worker_to_readers[worker] = std::vector<std::unique_ptr<TopicRecordReader>>();
          }
          worker_to_readers[worker].push_back(std::move(readers[i]));
        }

        std::vector<std::unique_ptr<orio::record::MergedTopicReader>> merge_readers;
        for (auto& [_worker, _readers] : worker_to_readers) {
          buffers.clear();
          for (int i = 0; i < _readers.size(); ++i) {
            buffers_.push_back(std::make_unique<char[]>(max_record_size));
            buffers.push_back(reinterpret_cast<orio::record::TopicRecord*>(buffers_.back().get()));
          }
          merge_readers.push_back(std::make_unique<orio::record::MergedTopicReader>(
              std::move(_readers), std::move(buffers)));
        }

        buffers.clear();
        for (int i = 0; i < merge_readers.size(); ++i) {
          buffers_.push_back(std::make_unique<char[]>(max_record_size));
          buffers.push_back(reinterpret_cast<orio::record::TopicRecord*>(buffers_.back().get()));
        }
        reader_ = std::make_unique<ShiftTopicReader>(std::move(merge_readers),
                                                     std::move(buffers));
      } else {
        reader_ = std::make_unique<orio::record::MergedTopicReader>(
            std::move(readers), std::move(buffers));
      }
    }
  }

  bool Read(exchange::base::executor::TopicRecordPacket* out) override {
    auto* record = reinterpret_cast<orio::record::TopicRecord*>(buf_.get());
    if (!reader_ || !reader_->Read(record)) {
      return false;
    }

    if (record->data_len > max_record_size_) {
      LOG(ERROR) << "abnormal size on FlowReader: " << record->data_len;
      return false;
    }

    out->sequence = record->topic_seq;
    out->payload = std::string(reinterpret_cast<char*>(record->data),
                               reinterpret_cast<char*>(record->data) + record->data_len);
    out->timestamp = record->timestamp;
    out->topic_name = topics_[record->topic_id];
    return true;
  }

  static std::vector<FlowFile> GetFileList(const std::string& machine, const std::string& root_dir,
                                           int64_t begin_timestamp,
                                           int64_t end_timestamp /* exclusive */) {
    std::vector<FlowFile> out;
    std::string last_path;

    for (int64_t ts = begin_timestamp; ts < end_timestamp; ts += 1 * 30 * 60 * 1000000000ULL) {
      char yyyymmdd[9];
      time_t t = ts / 1'000'000'000ULL;
      tm tm_{};
      gmtime_r(&t, &tm_);
      std::strftime(yyyymmdd, 9, "%Y%m%d", &tm_);
      auto path = fmt::format("{}/{}/{}/", root_dir, machine, yyyymmdd);
      if (path == last_path) {
        continue;
      } else {
        last_path = path;
      }

      std::vector<std::string> sub_paths;
      sub_paths.push_back(path);  // Add current dir first.
      for (const auto& entry : std::experimental::filesystem::directory_iterator(path)) {
        if (std::experimental::filesystem::is_directory(entry)) {
          sub_paths.push_back(entry.path().string());
        }
      }

      for (const auto& sub_path : sub_paths) {
        for (const auto& entry : std::experimental::filesystem::directory_iterator(sub_path)) {
          if (std::experimental::filesystem::is_directory(entry)) {
            continue;
          }

          // e.g. binance_rest_depth_2.queue-0030ai8kg8.20190203-000000Z.gz
          std::string filename = entry.path().filename();
          std::istringstream iss(filename);
          std::string queue;
          std::string time;
          FlowFile file;
          file.path = entry.path();

          if (!(getline(iss, file.topic_string, '.') && getline(iss, queue, '.') &&
                getline(iss, time, '.')) ||
              time.size() != 16) {
            throw std::runtime_error("ill-formatted filename: " + file.path.string());
          }

          tm tm2{};
          if (strptime(time.c_str(), "%Y%m%d-%H%M%S%Z", &tm2) == nullptr) {
            throw std::runtime_error("ill-formatted filename: " + file.path.string());
          }
          t = timegm(&tm2);
          // record file write to disk every 30 mins
          if ((t * 1'000'000'000LL < end_timestamp) &&
              ((t + 30 * 60) * 1'000'000'000LL >= begin_timestamp)) {
            file.timestamp = static_cast<int64_t>(t * 1'000'000'000LL);
            out.push_back(file);
          }
        }
      }
    }

    std::sort(out.begin(), out.end());
    return out;
  }

 private:
  const size_t max_record_size_;
  std::unique_ptr<char[]> buf_;
  std::unique_ptr<TopicRecordReader> reader_;
  std::vector<std::unique_ptr<char[]>> buffers_;
  std::vector<std::string> topics_;
  bool shift_;
};

}  // namespace coin2::flow
